/**
 * 引用自https://www.jianshu.com/p/f28dfbb2faa2
 * 2020/9/17
 */
import java.lang.instrument.ClassFileTransformer;
import java.security.ProtectionDomain;

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.CtNewMethod;

/**
 * 检测方法的执行时间
 */
public class MyTransformer implements ClassFileTransformer {

    final static String prefix = "\nlong startTime = System.currentTimeMillis();\n";
    final static String postfix = "\nlong endTime = System.currentTimeMillis();\n";

    /**
     * "the transformer will be called for every new class definition and every class redefinition."
     * 每个新类的定义和重定义都会调用transform方法，这里的指的是用ClassLoader.defineClass实现的定义，
     * 也就是说每个新类被ClassLoader加载的时候都会调用transform方法
     */
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
                            ProtectionDomain protectionDomain, byte[] classfileBuffer){
        // 仅处理测试类
        if(!className.startsWith("TestClass")){
            return null;
        }
        className = className.replace("/", ".");
        CtClass ctclass;
        try {
            // 使用全称,用于取得字节码类<使用javassist>
            ctclass = ClassPool.getDefault().get(className);
            for(CtMethod ctMethod : ctclass.getDeclaredMethods()){
                String methodName = ctMethod.getName();
                // 新定义一个方法叫做比如new$main
                String newMethodName = "new$" + methodName;
                // 将原来的方法名字修改
                ctMethod.setName(newMethodName);
                // 创建新的方法，复制原来的方法，名字为原来的名字
                CtMethod newMethod = CtNewMethod.copy(ctMethod, methodName, ctclass, null);
                // 构建新的方法体
                StringBuilder bodyStr = new StringBuilder();
                bodyStr.append("{");
                bodyStr.append("System.out.println(\"==============Enter Method: " + className + "." + methodName + " ==============\");");
                bodyStr.append(prefix);
                // 调用原有代码，类似于method();($$)表示所有的参数
                bodyStr.append(newMethodName + "($$);\n");
                bodyStr.append(postfix);
                bodyStr.append("System.out.println(\"==============Exit Method: " + className + "." + methodName + " Cost:\" +(endTime - startTime) +\"ms " + "===\");");
                bodyStr.append("}");

                // 替换新方法
                newMethod.setBody(bodyStr.toString());
                // 增加新方法
                ctclass.addMethod(newMethod);
            }
            // 修改后的方法列表 会发现多了一个方法
            System.out.println("after update method list ...");
            for(CtMethod ctMethod : ctclass.getDeclaredMethods()) {
                System.out.println(ctMethod.getName());
            }
            return ctclass.toBytecode();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }
}
